from libs.datasetFuncs import getDatas, getTargets
import numpy as np
from PIL import Image

class KNNClassifierPlus:
    def __init__(self, datasets, k=5):
        datas = getDatas(datasets, dtype='float32')
        self.datas = datas / 255
        self.targets = getTargets(datasets)
        self.k = k
    
    def predictProb(self, data):
        if isinstance(data, Image.Image):
            data = np.array(data, dtype='float32') / 255
        else:
            data = data.astype('float32') / 255
        distans = np.mean(np.abs(self.datas - data), axis=(1, 2))
        sorted_distans = np.sort(distans)
        limit = sorted_distans[self.k - 1]
        
        weights = 1 - distans
        weights = weights * ((weights >= 0.99) * 1 + (weights < 0.99) / (self.k - 1))
        predict_prob = np.zeros((np.unique(self.targets).size, ), dtype='float32')
        for i in np.where(distans <= limit)[0]:
            predict_prob[self.targets[i]] += weights[i]
        predict_prob /= predict_prob.sum()
        return predict_prob
    
    def predict(self, data):
        prob = self.predictProb(data)
        target = np.dot((prob == prob.max()), np.arange(prob.size))
        return target
